<html>
<head>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest"> </script>
</head>
<body>
    <button onclick="trainModel()">开始训练</button>
    <div id="output"></div>

    <script>
        // 1. 定义模型结构
        function createModel() {
            const model = tf.sequential();
            // 添加一个全连接层，输入维度为1，输出维度为1
            model.add(tf.layers.dense({ units: 1, inputShape: [1] }));
            return model;
        }

        // 2. 准备一些模拟数据 (y = 2x + 1)
        const xs = tf.tensor1d([1, 2, 3, 4, 5]);
        const ys = tf.tensor1d([3, 5, 7, 9, 11]); // 根据 y=2x+1 计算得出

        // 3. 创建模型、定义损失函数和优化器
        const model = createModel();
        model.compile({
            optimizer: 'sgd', // 使用随机梯度下降优化器
            loss: 'meanSquaredError' // 使用均方误差作为损失函数
        });

        // 4. 训练模型
        async function trainModel() {
            // 在UI上显示训练进度
            const outputDiv = document.getElementById('output');
            outputDiv.innerHTML = '训练中...';

            // 执行训练，迭代100次
            await model.fit(xs, ys, {
                epochs: 100,
                callbacks: {
                    onEpochEnd: async (epoch, logs) => {
                        outputDiv.innerHTML = `迭代次数: ${epoch+1}, 当前误差: ${logs.loss.toFixed(4)}`;
                        console.log(`Epoch ${epoch + 1}: loss = ${logs.loss}`);
                        // 为了不阻塞UI，每训练一步等待一下
                        await tf.nextFrame();
                    }
                }
            });

            // 5. 使用训练好的模型进行预测
            outputDiv.innerHTML += '<br>训练完成！<br>';
            const testX = tf.tensor1d([10]);
            const prediction = model.predict(testX);
            prediction.print();
            outputDiv.innerHTML += `对于输入 10，预测结果是：${prediction.dataSync()[0].toFixed(2)} (真实值应为 21)`;
        }
    </script>
</body>
</html>